#!/usr/bin/env python
# -*- conding: utf-8 -*-

"""
@Time     : 2024/8/12 21:47
@Author   : liujingmao
@File     : 2. StructureTool类方法.py
"""

from langchain_core.pydantic_v1 import BaseModel, Field
from langchain_core.tools import StructuredTool


class CaculatorInput(BaseModel):
    a: int = Field(..., description='第一个数字')
    b: int = Field(..., description='第二个数字')


def multiply(a: int, b: int) -> int:
    """convert two numbers to a number"""
    return a * b


async def amultiply(a: int, b: int) -> int:
    """convert two numbers to integer number"""
    return a * b


calculator = StructuredTool.from_function(
    func=multiply,
    coroutine=amultiply,
    name='multiply_tool',
    description="用于计算两个数的乘积",
    return_direct=True,
    args_schema=CaculatorInput,
)

# print the info of this tool
print("name:", calculator.name)
print("description:", calculator.description)
print("args:", calculator.args)
print("是否返回：", calculator.return_direct)

# call the tool
print(calculator.invoke({"a": 2, "b": 3}))
